#include "katanafilm.h"


#include <iostream>
#include <errno.h>
#include <netdb.h>
#include <string.h>
#include <unistd.h>
#include <unistd.h>

std::map<int, FnKat::NewFrameMessage* > KatanaFilm::_frameMap;
std::map<FnKat::NewFrameMessage*, int > KatanaFilm::_channelMap;

KatanaFilm::KatanaFilm(int xres, int yres, Filter *filt, const float crop[4],
                       bool openWindow, const string outputAddress,
                       int frameId, const float frameTime, string frameName) :
    ImageFilm( xres, yres, filt, crop, "", openWindow),
    outputAddress(outputAddress),
    frameId(frameId),
    frameTime(frameTime),
    frameName(frameName),
    isConnected(false),
    numChannels(3)
{
    bucketMutex = Mutex::Create();
}

KatanaFilm::~KatanaFilm(void)
{
    Mutex::Destroy(bucketMutex);
}

void KatanaFilm::UpdateDisplay(int x0, int y0, int x1, int y1, float splatScale)
{
    MutexLock lock(*bucketMutex);

    if(!isConnected)
    {
        isConnected = openConnection();
    }

    if (isConnected)
    {
        x0 = (x0 < 0)? 0: x0;
        y0 = (y0 < 0)? 0: y0;
        x1 = (x1 >= xResolution)? xResolution - 1: x1;
        y1 = (y1 >= yResolution)? yResolution - 1: y1;

        sendBucket(x0, y0, x1, y1, splatScale);
    }
}


void KatanaFilm::WriteImage(float splatScale)
{
    MutexLock lock(*bucketMutex);

    if (isConnected)
    {
        //sendBucket(0, 0, xResolution, yResolution, splatScale);
        closeConnection();
    }

    delete pipe;
    delete channel;
}


bool KatanaFilm::openConnection(void)
{
    // Parse the host:port string
    char hostName[4096];
    unsigned long portNumber;
    if (!parseHostname(outputAddress.c_str(), hostName, &portNumber))
        return false;

    pipe = FnKat::PipeSingleton::Instance(hostName, portNumber);

    // Try to connect pipe
    if (pipe->connect())
        return false;

    // First thing we need to do is check if we have a unique frame ID for
    // this render (remember a frame is made up of many channels)
    FnKat::NewFrameMessage *frameObject = NULL;
    if ( hasObjectForFrame(frameId) )
    {
        frameObject = getFrameObject(frameId);
    }
    else
    {
        // We are going to build one...
        frameObject = new FnKat::NewFrameMessage();
        frameObject->setFrameTime(frameTime);
        frameObject->setFrameDimensions(xResolution, yResolution);
        frameObject->setFrameOrigin(0, 0);

        std::string legacyEncoding;
        FnKat::encodeLegacyName(frameName, frameId, legacyEncoding);
        frameObject->setFrameName(legacyEncoding);

        // Store the frame in the map...
        addNewFrameObject(frameId, frameObject);

        // Send the frame down the pipe...
        pipe->send(*frameObject);
    }

    // Now we need to create a channel that will hold the data...
    int channelID = getNextChannelID(frameObject);
    FnKat::NewChannelMessage *channelObject = new FnKat::NewChannelMessage(*frameObject);

    // Sample rate (1.0, 1.0)
    float sampleRate[2];
    sampleRate[0] = 1.0;
    sampleRate[1] = 1.0;

    // Set channelID
    channelObject->setChannelID((FnKat::uint16_t)channelID);
    channelObject->setChannelOrigin(0, 0);
    channelObject->setChannelDimensions(xResolution, yResolution);
    channelObject->setSampleRate(sampleRate);
    channelObject->setDataSize(3 * sizeof(float));

    std::string legacyEncoding;
    FnKat::encodeLegacyName(frameName, frameId, legacyEncoding);
    channelObject->setChannelName(legacyEncoding);

    // Pack the channel down the pipe to and save to our image structure.
    pipe->send(*channelObject);
    channel = channelObject;

    return true;
}

void KatanaFilm::closeConnection(void)
{
    if (isConnected)
    {
        pipe->flushPipe(*channel);
        pipe->closeChannel(*channel);
    }
}

void KatanaFilm::sendBucket(int x0, int y0, int x1, int y1, float splatScale)
{
    int blockW = x1 - x0;
    int blockH = y1 - y0;
    size_t pixelSize = static_cast<size_t>(3 * sizeof(float));
    size_t rgbSize = static_cast<size_t>(blockW * blockH * pixelSize);

    // Convert image to RGB and compute final pixel values
    size_t offset = 0;
    float *rgb = new float[rgbSize];
    for (int y = y0; y < y1; ++y) {
        for (int x = x0; x < x1; ++x) {
            // Convert pixel XYZ color to RGB
            XYZToRGB((*pixels)(x, y).Lxyz, &rgb[offset]);

            // Normalize pixel with weight sum
            float weightSum = (*pixels)(x, y).weightSum;
            if (weightSum != 0.f) {
                float invWt = 1.f / weightSum;
                rgb[offset  ] = max(0.f, rgb[offset  ] * invWt);
                rgb[offset+1] = max(0.f, rgb[offset+1] * invWt);
                rgb[offset+2] = max(0.f, rgb[offset+2] * invWt);
            }

            // Add splat value at pixel
            float splatRGB[3];
            XYZToRGB((*pixels)(x, y).splatXYZ, splatRGB);
            rgb[offset++] += splatScale * splatRGB[0];
            rgb[offset++] += splatScale * splatRGB[1];
            rgb[offset++] += splatScale * splatRGB[2];
        }
    }

    // Init the Data Message
    FnKat::DataMessage *dataMsg = new FnKat::DataMessage(*channel);
    dataMsg->setStartCoordinates(x0, y0);
    dataMsg->setDataDimensions(blockW, blockH);

    dataMsg->setData(rgb, rgbSize);
    dataMsg->setByteSkip(pixelSize);

    // Pack it down the pipe.
    pipe->send(*dataMsg);

    // Clean up the data array we used to do the copy...
    delete [] rgb;
    rgb = NULL;

    delete dataMsg;
    dataMsg = NULL;
}


bool KatanaFilm::parseHostname(const char* hostnameString,
                               char* hostName,
                               unsigned long* portNumber)
{
    // Skip any leading _'s
    while (*hostnameString == '_') ++hostnameString;

    // Find the : after the hostname.
    const char* p = hostnameString;
    while ((p - hostnameString < 4096) && *p && ':' != *p)
    {
        ++p;
    }

    if (*p != ':')
    {
        return false;
    }

    // Copy out the hostname, null terminate it
    memcpy(hostName, hostnameString, p - hostnameString);
    hostName[p - hostnameString] = '\0';

    // Get the port number as a number.
    char* dummy;
    *portNumber = strtoul(p+1, &dummy, 0);

    *portNumber += 100;

    return true;
}


FnKat::NewFrameMessage* KatanaFilm::getFrameObject(int frameId)
{
    return _frameMap[frameId];
}


void KatanaFilm::addNewFrameObject(int frameId, FnKat::NewFrameMessage* frame)
{
  // Add the frame to the _frameMap...
  _frameMap[frameId] = frame;

  // Add the frame to the _channelMap.
  _channelMap[frame] = 0;
}


bool KatanaFilm::hasObjectForFrame(int frameId)
{
  std::map<int, FnKat::NewFrameMessage* >::iterator it;

  it = _frameMap.find(frameId);

  return (it != _frameMap.end());
}


int KatanaFilm::getNextChannelID(FnKat::NewFrameMessage *frame)
{
  // First find the std::set for the frame...
  std::map<FnKat::NewFrameMessage*, int>::iterator it = _channelMap.find(frame);

  // Return a new channel ID
  return ++(it->second);
}


//TODO: temporary hack to allow interacrive renders in Katana
KatanaFilm *CreateKatanaFilm(const ParamSet &params, Filter *filter)
{
    int xres = params.FindOneInt("xresolution", 640);
    int yres = params.FindOneInt("yresolution", 480);
    if (PbrtOptions.quickRender) xres = max(1, xres / 4);
    if (PbrtOptions.quickRender) yres = max(1, yres / 4);
    bool openwin = params.FindOneBool("display", false);
    float crop[4] = { 0, 1, 0, 1 };
    int cwi;
    const float *cr = params.FindFloat("cropwindow", &cwi);
    if (cr && cwi == 4) {
        crop[0] = Clamp(min(cr[0], cr[1]), 0., 1.);
        crop[1] = Clamp(max(cr[0], cr[1]), 0., 1.);
        crop[2] = Clamp(min(cr[2], cr[3]), 0., 1.);
        crop[3] = Clamp(max(cr[2], cr[3]), 0., 1.);
    }

    string outputAddress = params.FindOneString("outputAddress", "localhost:15000");
    int frameId = params.FindOneInt("frameId", 0);
    float frameTime = params.FindOneInt("frameTime", 0);
    string frameName = params.FindOneString("frameName", "Render");


    return new KatanaFilm(xres, yres, filter, crop, openwin,
                          outputAddress, frameId, frameTime, frameName);
}
